
#include "pubsub.hpp"
using namespace pubsub;

void State::add(const std::string& clientID, std::shared_ptr<Query> query,const CallbackFunction & callbackFunction)
{
    // add query if needed
    if(queries.count(query)==0)
    {
        queries[query] = std::map<std::string, CallbackFunction>();
    }

    //create subscription
    queries[query][clientID] = callbackFunction;

    //add client if needed
    if(clients.count(clientID)==0)
    {
        clients[clientID] = std::set<std::shared_ptr<Query>>();
    }
    clients[clientID].insert(query);
}

void State::remove(const std::string& clientID, std::shared_ptr<Query> query)
{
    if(queries.count(query)==0) return;
    auto clientFunctionMap= queries[query];

    // if the clientID is in the query's map
    auto fin=clientFunctionMap.find(clientID);
    if(fin!=clientFunctionMap.end())
    {
        clientFunctionMap.erase(fin);

        clients[clientID].erase(clients[clientID].find(query));

        if(clients[clientID].empty())
            clients.erase(clients.find(clientID));
    }
}

void State::removeAll(std::string clientID)
{
    if(clients.count(clientID) == 0) return;
    auto querySet=clients[clientID];

    for(auto item = querySet.begin();item != querySet.end(); item++)
    {
        auto clientFunctionMap = queries[(*item)];
        clientFunctionMap.erase(clientFunctionMap.find(clientID));
    }
    clients.erase(clients.find(clientID));
}


void State::call(const bytes& msg,const std::map<std::string, std::string>& tags)
{
    for(auto item=queries.begin();item!=queries.end();item++)
    {
        if(item->first->matches(tags))
        {
            for(auto fi=item->second.begin();fi!=item->second.end();fi++)
            {
                fi->second(msg);
            }
        }
    }
}


void Server::subscribe(const std::string& clientID, std::shared_ptr<Query> query,const CallbackFunction& cbFunction)
{
    std::unique_lock<std::mutex> lock(x_subscriptions_);
    if(subscriptions_.count(clientID)!=0)
    {
        if(subscriptions_[clientID].count(query->toString())!=0) return;
    } else
    {
        subscriptions_[clientID]=std::map<std::string, std::shared_ptr<Query>>();
    }
    {
        std::unique_lock<std::mutex> lock_cmd(x_cmds_);
        cmds_.push(cmd{
                .clientID=clientID,
                .op=sub,
                .query=query,
                .tags=std::map<std::string,std::string>(),
                .msg=std::vector<unsigned char>(),
                .callbackFunction=cbFunction
        });
    }
    condition_.notify_one();
    subscriptions_[clientID][query->toString()]=query;
}

void Server::unSubscribe(const std::string& clientID, std::shared_ptr<Query> query)
{
    std::unique_lock<std::mutex> lock(x_subscriptions_);
    auto origQuery = subscriptions_[clientID][query->toString()];
    {
        std::unique_lock<std::mutex> lock_cmd(x_cmds_);
        cmds_.push(cmd{
                .clientID=clientID,
                .op=unsub,
                .query=origQuery,
                .tags=std::map<std::string,std::string>(),
                .msg=std::vector<unsigned char>(),
                .callbackFunction=NULL
        });
    }
    condition_.notify_one();
    subscriptions_[clientID].erase(subscriptions_[clientID].find(query->toString()));
}

void Server::unSubscribeAll(const std::string& clientID)
{
    std::unique_lock<std::mutex> lock(x_subscriptions_);
    if(subscriptions_.count(clientID)==0)
        return;
    {
        std::unique_lock<std::mutex> lock_cmd(x_cmds_);
        cmds_.push(cmd{
                .clientID=clientID,
                .op=unsub,
                .query=std::shared_ptr<Query>(),
                .tags=std::map<std::string,std::string>(),
                .msg=std::vector<unsigned char>(),
                .callbackFunction=NULL
        });
    }
    condition_.notify_one();
    subscriptions_.erase(subscriptions_.find(clientID));
}

void Server::publich(const bytes& msg)
{
    std::map<std::string, std::string> tags;
    publichWithTags(msg,tags);
}

void Server::publichWithTags(const bytes &msg, const std::map<std::string, std::string> &tags)
{
    {
        std::unique_lock<std::mutex> lock_cmd(x_cmds_);
        cmds_.push(cmd{
                .clientID="",
                .op=pub,
                .query=std::shared_ptr<Query>(),
                .tags=tags,
                .msg=msg,
                .callbackFunction=NULL
        });
    }
    condition_.notify_one();
}

void Server::stop()
{
    {
        std::unique_lock<std::mutex> lock_cmd(x_cmds_);
        stop_=true;
    }
    condition_.notify_all();
}

Server::Server():worker_([this]()->void{
    for(;;)
    {
        cmd c;
        {
            std::unique_lock<std::mutex> lock(x_cmds_);
            this->condition_.wait(lock,[this](){return this->stop_ || !cmds_.empty();});
            if(this->stop_&&this->cmds_.empty())
                return;
            c=std::move(this->cmds_.front());
            this->cmds_.pop();
        }
        switch (c.op){
            case unsub:
                if(c.query.use_count()==0)  //remove all
                    state.removeAll(c.clientID);
                else
                    state.remove(c.clientID,c.query);
                break;
            case sub:
                state.add(c.clientID,c.query,c.callbackFunction);
                break;
            case pub:
                state.call(c.msg,c.tags);
                break;
        }
    }
}),stop_(false)
{
}

Server::~Server()
{
    {
        std::unique_lock<std::mutex> lock(x_cmds_);
        stop_=true;
    }
    condition_.notify_all();
    worker_.join();
}